import pymupdf
import argparse
import re
import cv2
import pytesseract
import numpy as np
from alibaba_get_table import submit_file, query, table_to_html

class PDF_ORC_HTML:
    def __init__(self):
        self.html = ''

    def list_to_html_table(self, lst):
        table_html = "<table>"
        # 添加表头
        table_html += "<thead><tr>"
        for item in lst[0]:
            table_html += f"<th>{item}</th>"
        table_html += "</tr></thead><tbody>"
        # 添加表格内容
        for row in lst[1:]:
            table_html += "<tr>"
            for item in row:
                table_html += f"<td>{item}</td>"
            table_html += "</tr>"
        table_html += "</tbody></table>"
        return table_html

    def text_box_judgment(self, tested_box, box_list):
        # 检测文本块是否在表格中,并添加表标识
        tested_box_x = float(tested_box[2] + tested_box[0]) / 2.0
        tested_box_y = float(tested_box[3] + tested_box[1]) / 2.0
        for index, box in enumerate(box_list):
            if box[2] > tested_box_x > box[0] and box[3] > tested_box_y > box[1]:
                return index + 1
        return 0

    def text_box_table_judgment(self, block):
        # 评估文本块是否为表格列表
        if 'image' in block:
            return 0
        block_text_list = []
        # [block[line[span]]]
        for line in block['lines']:
            x_ = []
            y1 = []
            dt = block['bbox'][2] - block['bbox'][0]
            a = ''
            b = 0
            for span in line['spans']:
                a += span['text']
                b += span['bbox'][2] - span['bbox'][0]
                y1.append(span['bbox'][1])
                x_.append([span['bbox'][0], span['bbox'][2]])
            block_text_list.append(a)
            odd = [o[-1] for o in x_[::2]]
            even = [e[0] for e in x_[1::2]]
            if len(block_text_list) >= 3 and dt - b > 5 and len(set([int(y) for y in y1])) == 1 and sum(
                    [y - x for x, y in zip(odd, even)]) > 10:
                return 99
        return 0

    def process_table(self, table):
        # 表格文本列表处理
        new_table = []
        for row in table:
            for index, j in enumerate(row):
                if j == None:
                    row[index] = ''
                if str('\n') in str(j):
                    row[index] = j.split('\n')
            if any(isinstance(x, list) for x in row if x):
                new_table += [list(r) for r in zip(*[x for x in row if x])]
                continue
            new_table.append([x for x in row])
        self.html += self.list_to_html_table(new_table)  # html处理

    def process_title():
        # 标题识别
        # lines:{'spans': [{'size': 18.0, 'flags': 4, 'font': 'SimHei', 'color': 0, 'ascender': 0.859375, 'descender': -0.140625,
        # 'text': '第1 章绪论', 'origin': (248.15899658203125, 101.80899810791016),
        # 'bbox': (248.15899658203125, 86.34024810791016, 347.15899658203125, 104.34024810791016)}],
        # 'wmode': 0, 'dir': (1.0, 0.0), 'bbox': (248.15899658203125, 86.34024810791016, 347.15899658203125, 104.34024810791016)}
        print('文本处理：标题识别+html转化')

    def image_tables_judging(self, image):
        # 判断图片是否存在表格:(暂时只能处理图片存在一个表的情况，后续多个表需要添加判断竖线来分割表格或设置横线距离进行剔除）
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)  # 灰度转化
        # 简单阈值进行图像二值化
        ret, thresh = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY)  #
        # 进行两次中值滤波
        blur = cv2.medianBlur(thresh, 3)  # 模板大小3*3
        blur = cv2.medianBlur(blur, 3)  # 模板大小3*3
        # 膨胀
        # dilation = cv2.dilate(blur,kernel,iterations = 1)
        h, w = gray.shape
        # 横向直线列表
        horizontal_lines = []
        for i in range(h - 1):
            # 找到两条记录的分隔线段，以相邻两行的平均像素差大于120为标准
            if abs(np.mean(blur[i, :]) - np.mean(blur[i + 1, :])) > 120:  # 255/2约为120，总长度大于页面一半判定为表格横线
                horizontal_lines.append([0, i, w, i])
        # print(horizontal_lines)
        return horizontal_lines

    def VerticalLineDetect(self, image):
        # 竖线检测,待完善
        # mask = 255 * np.ones(image.shape[:2], dtype = "uint8")  # 图片全选修复
        # image = cv2.inpaint(image, mask, 5, cv2.INPAINT_TELEA)  # 图像修复:cv2.INPAINT_NS、cv2.INPAINT_TELEA
        blur = cv2.GaussianBlur(image, (5, 5), 0)  # 高斯模糊
        gray = cv2.cvtColor(blur, cv2.COLOR_BGR2GRAY)  # 灰度转化
        edges = cv2.Canny(gray, 30, 240, apertureSize = 3)  # 二值图边缘检测
        minLineLength = 100
        maxLineGap = 10
        lines = cv2.HoughLinesP(edges, 1, np.pi / 180, 100, minLineLength, maxLineGap).tolist()
        for x1, y1, x2, y2 in lines[0]:
            cv2.line(image, (x1, y1), (x2, y2), (255, 0, 0), 5)
        cv2.imshow("houghline", image)
        cv2.waitKey()
        cv2.destroyAllWindows()
        print(lines)

    def image_ocr(self, image_path,alibaba_ocr=False):
        # 图片ocr识别 阿里图片表格识别api
        if alibaba_ocr:
            ID = submit_file(image_path, 'png')  # alibaba_api
            while ID:
                completed, json_str = query(ID)
                if completed:
                    if json_str['tables']:
                        self.html += table_to_html(json_str)
                    else:
                        return True
            return False
        # one_page = pymupdf.open('pdf', pdf_pixmap.pdfocr_tobytes(language = 'chi_sim',
        #                                                          tessdata = 'C:/Program Files/Tesseract-OCR/tessdata'))
        pytesseract.pytesseract.tesseract_cmd = r'.\Tesseract-OCR\tesseract.exe'
        image = cv2.imread(image_path)
        horizontal_lines = self.image_tables_judging(image)
        # 起码存在三条横线，判断存在表格
        if len(horizontal_lines) >= 6:
            # 有表格的图片处理
            # image_dst=cv2.fastNlMeansDenoisingColored(image,None,1,1,7,21)#高斯去噪
            mask = 255 * np.ones(image.shape[:2], dtype = "uint8")  # 图片全选修复
            image = cv2.inpaint(image, mask, 5, cv2.INPAINT_TELEA)  # 图像修复:cv2.INPAINT_NS、cv2.INPAINT_TELEA
            i_init = horizontal_lines[0][1]  # 指针
            table_text = ''
            text = pytesseract.image_to_string(image[:i_init, :], lang = 'chi_sim')
            self.html += '<p>' + text + '</p>'
            for i in horizontal_lines[::2][1:]:
                table_line = image[i_init:i[1], :]
                text = pytesseract.image_to_string(table_line, lang = 'chi_sim')
                table_text += text
                i_init = i[1]
            table_list = [(lambda s: [i for i in s.split() if i])(s) for s in table_text.split('\n') if s]  # text_str处理
            self.html += self.list_to_html_table(table_list)
            text = pytesseract.image_to_string(image[i_init:, :], lang = 'chi_sim')
            self.html += '<p>' + text + '</p>'
            return False
        else:
            # print('图片不存在表格，进行普通图片文本识别处理')
            return True

    def pdf_HTML(self,data_path,html_path,image_staging):
        with pymupdf.open(data_path) as pdf:
            tables_html = ''  # 存放表格html
            for page in pdf:
                # pdf_tables=pdf[i].find_tables(join_x_tolerance=5,snap_y_tolerance=4,text_x_tolerance=0.5,text_y_tolerance=0.5)#页面表格    (样例pdf第9页，标签不对齐识别出错）
                # 表格识别
                pdf_tables = page.find_tables(clip = (0, 50, 800, 750))
                tables_list = [pdf_tables[i].extract() for i in range(len(pdf_tables.tables))]  # 页面全部表格的文本列表文本
                tables_box = [pdf_tables[i].bbox for i in range(len(pdf_tables.tables))]  # 页面全部表格的方框列表，方框格式：元组
                blocks = page.get_text("dict")["blocks"]
                for i, block in enumerate(blocks):
                    blocks[i]['table-type'] = self.text_box_judgment(block['bbox'], tables_box)
                    if not blocks[i]['table-type']:
                        blocks[i]['table-type'] = self.text_box_table_judgment(block)  # 0/99
                last_table_type = 0  # 初始指针
                wire_zhizhen = 0
                wire_text = []  # 无线表格
                for i, block in enumerate(blocks):
                    if block['table-type'] == 99:
                        # print('无线表格处理')
                        a = []
                        for j in block['lines']:
                            a.append(''.join([k['text'] for k in j['spans']]))
                        wire_text.append(a)
                        wire_zhizhen += 1
                    elif last_table_type == 99 and wire_zhizhen >= 3:
                        self.process_table(wire_text)
                        wire_text = []
                    else:
                        for w in wire_text:
                            self.html += '<p>' + ''.join(w) + '</p>'
                        wire_zhizhen = 0
                        wire_text = []

                    if block['table-type'] == 0:
                        text = ''
                        if 'image' not in block:
                            # print('正常文本处理')
                            text += ''.join([l['text'] for k in block['lines'] for l in k['spans']])
                            self.html += '<p>' + text + '</p>'
                        else:
                            # print('图片处理')
                            pdf_pixmap = page.get_pixmap(dpi = 300, clip = block['bbox'])
                            pdf_pixmap.save(image_staging)  # 暂存png图片，可进行编号全部保存
                            if self.image_ocr(image_staging):
                                # 普通图片处理
                                one_page = pymupdf.open('pdf', pdf_pixmap.pdfocr_tobytes(language = 'chi_sim',
                                                                                         tessdata = './Tesseract-OCR/tessdata'))
                                for j in one_page[0].get_text("dict")["blocks"][1:]:
                                    text += ''.join([l['text'] for k in j['lines'] for l in k['spans']])
                                    for k in j['lines']:
                                        for l in k['spans']:
                                            text += l['text']
                                self.html += '<p>' + text + '</p>'
                    else:
                        if block['table-type'] > last_table_type and block['table-type'] != 99:
                            # print('正常表格处理')
                            self.process_table(tables_list[block['table-type'] - 1])
                    last_table_type = block['table-type']
        with open(html_path, 'w') as f:
            f.write(self.html)


if __name__=='__main__':
    parser=argparse.ArgumentParser()
    parser.add_argument("--pdf_path", type = str, default = "./data/test.pdf", help = "pdf文件路径")
    parser.add_argument("--html_save_path", type = str, default = './output/test.html', help = "html文件保存路径")
    parser.add_argument("--image_staging", type = str, default = './staging/test.png', help = "图片暂存")
    args = parser.parse_args()

    init = PDF_ORC_HTML()
    init.pdf_HTML(args.pdf_path, args.html_save_path,args.image_staging)